import torch
from torch import nn
from d2l import torch as d2l
#########################################################
#卷积
def corr2d(X,K):
    h,w = K.shape
    Y = torch.zeros((X.shape[0]-h+1,X.shape[1]-w+1))
    for i in range(Y.shape[0]):
        for j in range(Y.shape[1]):
            Y[i,j] = (X[i:i+h,j:j+w]*K).sum()
    return Y

def corr2d_multi_in(X,K):
    print("------------")
    for a,b in zip(X,K):
        print(a)
        print(b)
        print('\n')
    return sum(corr2d(x,k) for x,k in zip(X,K))

def corr2d_multi_in_out(X,K):
    return torch.stack([corr2d_multi_in(X,k) for k in K],0)

############################################################
#池化
def pool2d(X,poolsize,mode='max'):
    p_h,p_w = poolsize
    Y = torch.zeros((X.shape[0]-p_h+1,X.shape[1]-p_w+1))
    for i in range(X.shape[0]):
        for j in range(X.shape[1]):
            if mode=='max':
                Y[i,j] = X[i:i+p_h,j:j+p_w].max()
            elif mode == 'avg':
                Y[i,j] = X[i:i+p_h,j:j+p_w].mean()
    return Y


class Conv2D(nn.Module):
    def __init__(self,kernel_size):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(kernel_size))
        self.bias = nn.Parameter(torch.zeros(1))

    def forward(self,x):
        return corr2d(x,self.weight)+self.bias
    
X = torch.tensor([[[0.0,1.0,2.0],
                   [3.0,4.0,5.0],
                   [6.0,7.0,8.0]],
                  [[1.0,2.0,3.0],
                   [4.0,5.0,6.0],
                   [7.0,8.0,9.0]]])
K = torch.tensor([[[0.0,1.0],
                   [2.0,3.0]],
                  [[1.0,2.0],
                  [3.0,4.0]]])

K = torch.stack((K,K+1,K+2),0)
print('K:')
print(K)
print('output:')
print(corr2d_multi_in_out(X,K))
